import sys
import os
import datetime
import arcpy
from arcpy.sa import *

# Inputs
conveyance_path = sys.argv[1]
local_flow_path = sys.argv[2]
flow_dir_path = sys.argv[3]
threshold = int(sys.argv[4])
scratch_gdb = sys.argv[5]
log_path = sys.argv[6]

def log(message):
    time_stamp = datetime.datetime.now().strftime('%x %X')
    full_message = "{0} - {1}".format(time_stamp, message)
    print full_message
    if log_path.strip() != "":
        log_file = open(log_path, 'a')
        log_file.write(full_message + "\n")
        log_file.close()

def calc_reclass(minimum, maximum, increment, factor = 1.05):
    remap = []
    lower = int(increment)
    upper = int(lower + (increment * factor))
    while upper < maximum:
        remap.append([lower, upper, lower])
        increment = upper - lower
        lower = upper
        upper = int(lower + (increment * factor))
    if upper >= maximum:
        upper = int(maximum)
        remap.append([lower, upper, lower])
    return arcpy.sa.RemapRange(remap)

try:
    # Check out license
    log("Getting Spatial Analyst license...")
    arcpy.CheckOutExtension("Spatial")

    output_gdb = os.path.dirname(conveyance_path)
    flow_dir = arcpy.sa.Raster(flow_dir_path)

    arcpy.env.overwriteOutput = True
    arcpy.env.workspace = output_gdb
    arcpy.env.scratchWorkspace = scratch_gdb
    arcpy.env.snapRaster = flow_dir
    arcpy.env.extent = flow_dir.extent

    # Extract outlet flow lines
    log("Extracting outlet flow lines...")
    outlet_lines = arcpy.FeatureClassToFeatureClass_conversion(conveyance_path, output_gdb, "outlet_drain_network", '"Subsurface"=0')
    log("Outlet flow lines: {0}".format(outlet_lines))

    # Extract vertices
    log("Extracting vertices from outlet flow lines...")
    vertices = arcpy.FeatureVerticesToPoints_management(outlet_lines, "outlet_drain_network_vertices")
    log("Outlet flow line vertices: {0}".format(vertices))

    # Convert to raster
    log("Converting vertices to raster...")
    outlet_flow = arcpy.PointToRaster_conversion(vertices, "accumulated_area_sq_ft", "outlet_runoff_accumulation", "SUM", "", 3)
    log("Outlet flow grid: {0}".format(outlet_flow))

    #TODO: Set output raster extent/snap grid

    # Combining local and outlet flow
    log("Creating combined runoff accumulation raster...")
    combined_flow = Int(Con(IsNull(outlet_flow), 0, outlet_flow) + Raster(local_flow_path))
    log("Combined flow accumulation: {0}".format(combined_flow))

    # Extract channel flow by threshold
    log("Creating threshold runoff accumulation raster...")
    reclass_table = calc_reclass(combined_flow.minimum, combined_flow.maximum, threshold)
    threshold_flow = arcpy.sa.Reclassify(combined_flow, "VALUE", reclass_table, "NODATA")
    log("Threshold runoff accumulation: {0}".format(threshold_flow))

    # Convert to stream lines
    log("Creating stream lines from raster...")
    stream_lines = StreamToFeature(Int(threshold_flow), flow_dir, "accumulated_surface_lines", "NO_SIMPLIFY")
    log("Streamline feature class: {0}".format(stream_lines))

except arcpy.ExecuteError:
    log(arcpy.GetMessages(2))
    sys.exit()
except Exception as e:
    log("Error: {0}".format(e.message))
    sys.exit()
finally:
    log("Releasing Spatial Analyst license...")
    arcpy.CheckInExtension("Spatial")
    log("Total flow accumulation complete.")
